# 1. 安装依赖（requirements.txt需包含这些）
# transformers==4.36.2
# datasets==2.14.6
# peft==0.7.1
# accelerate==0.25.0
# bitsandbytes==0.41.1
# torch==2.1.0
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import json
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    BitsAndBytesConfig
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training


# 2. 配置参数（根据你的需求修改）
class Config:
    # 模型路径（本地Ollama的deepseek-r1:7b路径，或Hugging Face模型名）
    model_name_or_path = "mistralai/Mistral-7B-v0.1" # 建议用Hugging Face官方权重
    # 数据集路径
    data_path = "../data/finetune/planner_data.json"
    # 输出LoRA权重路径
    output_dir = "./lora_weights/planner"
    # 训练参数
    per_device_train_batch_size = 2  # 16GB显存用2，32GB可用4
    gradient_accumulation_steps = 4  # 梯度累积，模拟更大batch
    learning_rate = 2e-4  # LoRA微调常用学习率
    num_train_epochs = 3  # 小数据集3轮足够
    fp16 = True  # 混合精度训练，节省显存
    logging_steps = 10  # 每10步打印一次日志
    save_steps = 100  # 每100步保存一次权重


# 3. 加载并格式化数据集（规划Agent需要输出JSON格式的决策）
def load_dataset(data_path):
    with open(data_path, "r", encoding="utf-8") as f:
        data = json.load(f)  # 格式：[{"query": "...", "output": "..."}]

    # 格式化prompt（让模型学习“输入query→输出JSON”的映射）
    def format_prompt(example):
        prompt = f"""用户查询：{example['query']}
请分析是否需要调用本地RAG或MCP搜索，输出JSON（仅JSON，无多余内容）：
{example['output']}"""
        return {"text": prompt}

    # 转换为Hugging Face Dataset格式
    dataset = Dataset.from_list(data)
    dataset = dataset.map(format_prompt)
    return dataset


# 4. 加载模型并配置LoRA
def load_model_and_lora(config):
    # 直接删除原来的 bnb_config 代码块！

    # 4.2 加载底座模型（修改后：无量化，用 FP16）
    model = AutoModelForCausalLM.from_pretrained(
        config.model_name_or_path,
        torch_dtype=torch.float16,  # 用 FP16 精度（16GB显存够用）
        device_map="auto",  # 自动分配 GPU/CPU
        trust_remote_code=True
    )
    # 准备模型用于量化训练
    model = prepare_model_for_kbit_training(model)

    # 4.3 配置LoRA参数（针对规划Agent的任务特点）
    lora_config = LoraConfig(
        r=16,  # 秩，控制LoRA能力，16适合7B模型
        lora_alpha=32,
        # 只训练注意力层，DeepSeek的注意力层名是q_proj、k_proj、v_proj、o_proj
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )

    # 4.4 绑定LoRA到模型
    model = get_peft_model(model, lora_config)
    # 打印可训练参数比例（通常只有1-2%）
    model.print_trainable_parameters()
    return model


# 5. 配置Tokenizer（处理文本编码）
def load_tokenizer(model_name_or_path):
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    tokenizer.pad_token = tokenizer.eos_token  # 设置padding符号
    return tokenizer


# 6. 数据预处理（将文本转换为模型可接受的输入）
def preprocess_function(examples, tokenizer):
    # 用tokenizer编码文本，最长512 tokens
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,
        padding="max_length",
        return_tensors="pt"
    )


# 7. 主训练函数
def main():
    config = Config()

    # 加载数据集
    dataset = load_dataset(config.data_path)
    # 加载tokenizer
    tokenizer = load_tokenizer(config.model_name_or_path)
    # 预处理数据集
    tokenized_dataset = dataset.map(
        lambda x: preprocess_function(x, tokenizer),
        batched=True  # 批量处理
    )

    # 加载模型和LoRA
    model = load_model_and_lora(config)

    # 配置训练参数
    training_args = TrainingArguments(
        output_dir=config.output_dir,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        num_train_epochs=config.num_train_epochs,
        fp16=config.fp16,
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        optim="paged_adamw_8bit",  # 8bit优化器，更省显存
        report_to="none"  # 不使用wandb等日志工具
    )

    # 初始化Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset
    )

    # 开始训练
    print("开始训练规划Agent的LoRA...")
    trainer.train()

    # 保存最终LoRA权重
    model.save_pretrained(f"{config.output_dir}/final")
    print(f"LoRA权重已保存到：{config.output_dir}/final")


if __name__ == "__main__":
    main()
